#!/bin/bash

# 颜色定义
RED='\033[31m'
GREEN='\033[32m'
YELLOW='\033[33m'
BLUE='\033[34m'
NC='\033[0m' # 重置颜色

# 显示标题
echo -e "${GREEN}\n===== 存储设备I/O自动优化 =====${NC}"
echo -e "时间: $(date '+%Y-%m-%d %H:%M:%S')"
echo -e "内核版本: $(uname -r)\n"

# 获取所有物理块设备
PHYSICAL_DEVS=()
while IFS= read -r dev; do
    # 排除虚拟设备和只读设备
    if [[ ! $dev =~ ^(loop|ram|fd|dm-|md|sr) ]] && 
       [ $(cat /sys/block/${dev#/dev/}/ro) -eq 0 ]; then
        PHYSICAL_DEVS+=("$dev")
    fi
done < <(lsblk -d -no NAME | awk '{print "/dev/"$1}')

# 检查是否找到设备
if [ ${#PHYSICAL_DEVS[@]} -eq 0 ]; then
    echo -e "${RED}错误: 未找到可写的物理存储设备${NC}"
    exit 1
fi

# 显示检测到的设备
echo -e "${YELLOW}[检测到以下物理设备]${NC}"
for dev in "${PHYSICAL_DEVS[@]}"; do
    dev_name=${dev#/dev/}
    model=$(cat /sys/block/$dev_name/device/model 2>/dev/null | tr -d '\n' | sed 's/\s*$//')
    size=$(lsblk -d -no SIZE $dev)
    rotational=$(cat /sys/block/$dev_name/queue/rotational)
    type=$([ $rotational -eq 1 ] && echo "HDD" || echo "SSD/NVMe")
    scheduler=$(cat /sys/block/$dev_name/queue/scheduler | grep -o '\[.*\]' | tr -d '[]')
    
    printf "%-8s %-20s %-8s %-12s %-15s\n" \
        "$dev" "${model:-N/A}" "$size" "$type" "$scheduler"
done
echo

# 1. 设备级优化
for DEVICE in "${PHYSICAL_DEVS[@]}"; do
    DEVICE_NAME=${DEVICE#/dev/}
    ROTATIONAL=$(cat /sys/block/$DEVICE_NAME/queue/rotational)
    
    echo -e "${BLUE}优化设备: $DEVICE${NC}"
    
    # 预读设置 (HDD: 8MB, SSD: 2MB)
    RA_VALUE=$([ $ROTATIONAL -eq 1 ] && echo "16384" || echo "4096")
    blockdev --setra $RA_VALUE $DEVICE 2>/dev/null
    echo -e "├─ 预读: ${GREEN}$((RA_VALUE * 512 / 1024))KB${NC} (原: $(( $(blockdev --getra $DEVICE) * 512 / 1024 ))KB)"
    
    # I/O调度器优化
    SCHED_PATH="/sys/block/$DEVICE_NAME/queue/scheduler"
    CURRENT_SCHED=$(grep -o '\[.*\]' $SCHED_PATH | tr -d '[]')
    AVAILABLE_SCHEDS=$(cat $SCHED_PATH | sed 's/.*\[//; s/\].*//')
    
    if [ $ROTATIONAL -eq 1 ]; then
        # HDD优化
        TARGET_SCHED="mq-deadline"
        [[ ! $AVAILABLE_SCHEDS =~ $TARGET_SCHED ]] && TARGET_SCHED="bfq"
    else
        # SSD/NVMe优化
        TARGET_SCHED="none"
        [[ ! $AVAILABLE_SCHEDS =~ $TARGET_SCHED ]] && TARGET_SCHED="kyber"
    fi
    
    # 应用调度器
    if [[ $AVAILABLE_SCHEDS =~ $TARGET_SCHED ]]; then
        echo $TARGET_SCHED > $SCHED_PATH 2>/dev/null
        echo -e "├─ 调度器: ${GREEN}$TARGET_SCHED${NC} (原: $CURRENT_SCHED)"
    else
        echo -e "├─ 调度器: ${YELLOW}保持 $CURRENT_SCHED (无$TARGET_SCHED)${NC}"
    fi
    
    # 队列深度优化
    NR_REQUESTS_PATH="/sys/block/$DEVICE_NAME/queue/nr_requests"
    if [ -f $NR_REQUESTS_PATH ]; then
        OLD_DEPTH=$(cat $NR_REQUESTS_PATH)
        NEW_DEPTH=$([ $ROTATIONAL -eq 1 ] && echo "256" || echo "1024")
        echo $NEW_DEPTH > $NR_REQUESTS_PATH 2>/dev/null
        echo -e "└─ 队列深度: ${GREEN}$NEW_DEPTH${NC} (原: $OLD_DEPTH)"
    fi
    echo
done

# 2. 全局虚拟内存优化
echo -e "${YELLOW}[全局虚拟内存优化]${NC}"
sysctl -w vm.dirty_ratio=10 2>/dev/null
sysctl -w vm.dirty_background_ratio=5 2>/dev/null
sysctl -w vm.dirty_expire_centisecs=3000 2>/dev/null
sysctl -w vm.swappiness=10 2>/dev/null

echo -e "├─ vm.dirty_ratio: ${GREEN}10${NC} (默认: $(sysctl -n vm.dirty_ratio 2>/dev/null || echo "?"))"
echo -e "├─ vm.dirty_background_ratio: ${GREEN}5${NC} (默认: $(sysctl -n vm.dirty_background_ratio 2>/dev/null || echo "?"))"
echo -e "├─ vm.dirty_expire_centisecs: ${GREEN}3000${NC} (默认: $(sysctl -n vm.dirty_expire_centisecs 2>/dev/null || echo "?"))"
echo -e "└─ vm.swappiness: ${GREEN}10${NC} (默认: $(sysctl -n vm.swappiness 2>/dev/null || echo "?"))"


# 完成提示
echo -e "\n${GREEN}优化完成! 建议重启后观察效果${NC}"
# echo "echo never > /sys/kernel/mm/transparent_hugepage/enabled"
echo "-------------------------------------------------"